from twisted.internet.selectreactor import SelectReactor
from twisted.python import context


class AsyncContextableReactor(SelectReactor):
    """A customized reactor supporting asynchronous context.

    After wrapping your code with twisted.python.context.call, all future code -
    no matter from callLater, callFromThread, callInThread (yes, threadpool by
    default support it) or incoming asynchronous calls - would be able to get
    the same context through twisted.python.context.get.

    See the test code for more usage in examples.
    """
    def currentContext(self):
        return context.theContextTracker.currentContext().contexts[-1]

    def addReader(self, reader):
        ctx = self.currentContext()
        self._reads[reader] = ctx

    def addWriter(self, writer):
        ctx = self.currentContext()
        self._writes[writer] = ctx

    def _doReadOrWrite(self, selectable, method, dict):
        ctx = (method == 'doRead' and self._reads or self._writes)[selectable]
        context.call(ctx, SelectReactor._doReadOrWrite, self, selectable,
                     method, dict)

    def callLater(self, _seconds, _f, *args, **kw):
        ctx = self.currentContext()
        return SelectReactor.callLater(self, _seconds, context.call, ctx, _f,
                                       *args, **kw)

    def callFromThread(self, f, *args, **kw):
        ctx = self.currentContext()
        return SelectReactor.callFromThread(self, context.call, ctx, f, *args,
                                            **kw)


def install():
    reactor = AsyncContextableReactor()
    from twisted.internet.main import installReactor
    installReactor(reactor)
